import json
from dataclasses import dataclass
from typing import List


# 定义 batch_task 类
@dataclass
class BatchTask:
    def __init__(self,
                 userId: int = None,
                 taskId: int = None,
                 modelPath: str = None,
                 imgPaths: List[str] = None,
                 imgsIds: List[int] = None,
                 savePath: str = None,
                 backClass: str = None,
                 userRoot: str = None,
                 labelSetId: int = None,):

        self.userId = userId
        self.taskId = taskId
        self.modelPath = modelPath
        self.imgPaths = imgPaths
        self.imgsIds = imgsIds
        self.savePath = savePath
        self.backClass = backClass
        self.userRoot = userRoot
        self.labelSetId = labelSetId

    def load_data(self, json_Str):
        args = json.loads(json_Str)
        self.__dict__ = args
        return self

    def __repr__(self):
        return str(self.__dict__)


# 定义train_task类
@dataclass
class TrainTask:
    def __init__(self,
                 userId: int = None,
                 taskId: int = None,
                 imgPaths: List[str] = None,
                 labelPaths: List[str] = None,
                 savePath: str = None,
                 labelClass: str = None,
                 image_type: str = 'CT'):
        self.userId = userId
        self.taskId = taskId
        self.imgPaths = imgPaths
        self.labelPaths = labelPaths
        self.savePath = savePath
        self.labelClass = labelClass
        self.image_type = image_type

    def load_data(self, json_Str):
        args = json.loads(json_Str)
        self.__dict__ = args
        return self

    def __repr__(self):
        return str(self.__dict__)


@dataclass
class UpdateModel:
    def __init__(self, userId=None, niiPath=None, savePath=None):
        self.userId = userId
        self.niiPath = niiPath
        self.savePath = savePath

    def load_data(self, json_str):
        args = json.loads(json_str)
        self.__dict__ = args
        return self
